// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <pollux/common/caching/file_ids.h>
#include <pollux/common/caching/scan_tracker.h>
#include <pollux/common/io/io_statistics.h>
#include <pollux/dwio/common/input_stream.h>
#include <pollux/dwio/common/seekable_input_stream.h>

namespace kumo::pollux::dwio::common {
    class DirectBufferedInput;

    /// An input stream over possibly coalesced loads. Created by
    /// DirectBufferedInput. Similar to CacheInputStream but does not use cache.
    class DirectInputStream : public SeekableInputStream {
    public:
        DirectInputStream(
            DirectBufferedInput *bufferedInput,
            IoStatistics *ioStats,
            const pollux::common::Region &region,
            std::shared_ptr<ReadFileInputStream> input,
            uint64_t fileNum,
            std::shared_ptr<cache::ScanTracker> tracker,
            cache::TrackingId trackingId,
            uint64_t groupId,
            int32_t loadQuantum);

        bool Next(const void **data, int *size) override;

        void BackUp(int count) override;

        bool SkipInt64(int64_t count) override;

        google::protobuf::int64 ByteCount() const override;

        void seekToPosition(PositionProvider &position) override;

        std::string getName() const override;

        size_t positionSize() const override;

        /// Testing function to access loaded state.
        void testingData(
            pollux::common::Region &loadedRegion,
            memory::Allocation *&data,
            std::string *&tinyData) {
            loadedRegion = loadedRegion_;
            data = &data_;
            tinyData = &tinyData_;
        }

    private:
        // Ensures that the current position is covered by 'data_'.
        void loadPosition();

        // Synchronously sets 'data_' to cover loadedRegion_'.
        void loadSync();

        DirectBufferedInput *const bufferedInput_;
        IoStatistics *const ioStats_;
        const std::shared_ptr<ReadFileInputStream> input_;
        // The region of 'input' 'this' ranges over.
        const pollux::common::Region region_;
        const uint64_t fileNum_;
        std::shared_ptr<cache::ScanTracker> tracker_;
        const cache::TrackingId trackingId_;
        const uint64_t groupId_;

        // Maximum number of bytes read from 'input' at a time.
        const int32_t loadQuantum_;

        // The part of 'region_' that is loaded into 'data_'/'tinyData_'. Relative to
        // file start.
        pollux::common::Region loadedRegion_;

        // Allocation with loaded data. Has space for region.length or loadQuantum_
        // bytes, whichever is less.
        memory::Allocation data_;

        // Contains the data if the range is too small for Allocation.
        std::string tinyData_;

        // Pointer  to start of current run in 'entry->data()' or
        // 'entry->tinyData()'.
        uint8_t *run_{nullptr};

        // Offset of current run from start of 'data_'
        uint64_t offsetOfRun_;

        // Position of stream relative to 'run_'.
        int offsetInRun_{0};

        // Index of run in 'data_'
        int runIndex_ = -1;

        // Number of valid bytes starting at 'run_'
        uint32_t runSize_ = 0;
        // Position relative to 'region_.offset'.
        uint64_t offsetInRegion_ = 0;

        // Set to true when data is first loaded.
        bool loaded_{false};
    };
} // namespace kumo::pollux::dwio::common
